from mcp.server.fastmcp import FastMCP, Context
import subprocess
import json
from typing import Optional

mcp = FastMCP("Ionice MCP Server")

"""
面向 ionice 软件的 MCP Server

MCP tools 列表：
1. ionice_query_priority - 查询进程的IO调度优先级
   - 用途：查询指定进程ID的当前IO调度优先级
   - 参数：pid - 进程ID
   - 返回：包含命令执行结果的JSON结构

2. ionice_set_priority - 设置进程的IO调度优先级
   - 用途：设置指定进程的IO调度类和优先级
   - 参数：pid - 进程ID，class_level - IO调度类（0-3），priority_level - 优先级（0-7）
   - 返回：包含命令执行结果的JSON结构
"""

@mcp.tool()
def ionice_query_priority(pid: int) -> dict:
    """
    查询指定进程的IO调度优先级
    
    参数:
        pid: 进程ID
        
    返回JSON结构:
        {
            "success": bool,      # 命令是否成功执行
            "command": str,       # 执行的完整命令
            "exit_code": int,     # 命令退出码
            "stdout": str,        # 标准输出
            "stderr": str         # 标准错误
        }
    """
    command = ["ionice", "-p", str(pid)]
    
    try:
        result = subprocess.run(command, capture_output=True, text=True, timeout=30)
        
        return {
            "success": result.returncode == 0,
            "command": " ".join(command),
            "exit_code": result.returncode,
            "stdout": result.stdout,
            "stderr": result.stderr
        }
    except subprocess.TimeoutExpired:
        return {
            "success": False,
            "command": " ".join(command),
            "exit_code": -1,
            "stdout": "",
            "stderr": "Command timed out"
        }
    except Exception as e:
        return {
            "success": False,
            "command": " ".join(command),
            "exit_code": -1,
            "stdout": "",
            "stderr": str(e)
        }

@mcp.tool()
def ionice_set_priority(pid: int, class_level: int, priority_level: int) -> dict:
    """
    设置指定进程的IO调度优先级
    
    参数:
        pid: 进程ID
        class_level: IO调度类 (0: none, 1: realtime, 2: best-effort, 3: idle)
        priority_level: 优先级 (0-7, 0为最高优先级)
        
    返回JSON结构:
        {
            "success": bool,      # 命令是否成功执行
            "command": str,       # 执行的完整命令
            "exit_code": int,     # 命令退出码
            "stdout": str,        # 标准输出
            "stderr": str         # 标准错误
        }
    """
    command = ["ionice", "-c", str(class_level), "-n", str(priority_level), "-p", str(pid)]
    
    try:
        result = subprocess.run(command, capture_output=True, text=True, timeout=30)
        
        return {
            "success": result.returncode == 0,
            "command": " ".join(command),
            "exit_code": result.returncode,
            "stdout": result.stdout,
            "stderr": result.stderr
        }
    except subprocess.TimeoutExpired:
        return {
            "success": False,
            "command": " ".join(command),
            "exit_code": -1,
            "stdout": "",
            "stderr": "Command timed out"
        }
    except Exception as e:
        return {
            "success": False,
            "command": " ".join(command),
            "exit_code": -1,
            "stdout": "",
            "stderr": str(e)
        }

if __name__ == "__main__":
    mcp.run()